# -*- coding: utf-8 -*-
from gwac_common import *
import numpy as np
import os
import sys
import glob
import shutil
from datetime import datetime 
####
start_time = datetime.now()
print('{0} Start'.format(start_time))
files={}
raindex=-1
decindex=-1
magindex=-1
####

#追加内容到文件
def appendFile(file, line):
    global files
    f = files.get(file)
    if f is None:
        makeParentDirs(file)
        f = open(file, 'w')
        files[file]=f
    f.write(line)

def closeFiles(files):
    for fname in files:
        files[fname].close()
    files={}

def detectColIndex(arr, col):
    count = len(arr)
    for i in range(0, count):
        if col==arr[i]:return i
    return -1


'''
在一台机器上做好分割，再分别复制到各个工作站上
规则为 rootDir/datasetName/dec/ra_mag.csv
'''

def splitFile2Dec(catalog, rootDir, datasetName, skipFirstLine=False):
    if not os.path.exists(catalog):
        return
    with open(catalog, 'r') as f:
        if skipFirstLine:
            line=f.readline()#pass header
        global raindex
        global decindex
        global magindex
        for line in f:
            arr=line.split(',')
            ra=float(arr[raindex])
            dec=float(arr[decindex])
            mag=float(arr[magindex])
            destfile="{0}/{1}/{2}.csv".format(rootDir, datasetName, int(np.floor(dec)))
            appendFile(destfile, line)

def splitFile2RaMag(catalog, rootDir, datasetName):
    if not os.path.exists(catalog):
        return
    with open(catalog, 'r') as f:
        global raindex
        global decindex
        global magindex
        for line in f:
            arr=line.split(',')
            ra=float(arr[raindex])
            dec=float(arr[decindex])
            mag=float(arr[magindex])
            destfile="{0}/{1}/{2}/{3}_{4}.csv".format(rootDir, datasetName, int(np.floor(dec)), int(np.floor(ra)), getMagId(mag))
            appendFile(destfile, line)
#
def testSplit():
    srcCatalog='/data/gaia_dr2/gaia2-j2000-for-gwac.csv'
    rootDir='/data/gaia_dr2/split-catalog/'
    datasetName='gaia_dr2'
    #
    global files
    global raindex
    global decindex
    global magindex
    arr='source_id,healpixid,htmid,ra,dec,pmra,pmdec,parallax,g_mag,bp_mag,rp_mag,teff,ra2015,dec2015,parallax2015,pmra2015,pmdec2015'.split(',')
    raindex=detectColIndex(arr, 'ra')
    decindex=detectColIndex(arr, 'dec')
    magindex=detectColIndex(arr, 'g_mag')
    print('{0} splitting catalog to dec'.format(datetime.now()))
    ##splitFile2Dec(srcCatalog, rootDir, datasetName, skipFirstLine=True)
    ##closeFiles(files)
    files={}
    ##
    print('{0} splitting catalog to ra_mag for each dec'.format(datetime.now()))
    for dec in range(-90, 90):
        catalogDir="{0}/{1}/{2}".format(rootDir, datasetName, int(np.floor(dec)))
        if os.path.exists(catalogDir):
            print('skip dec=',dec)
            continue
        else:
            print('handling dec=', dec)
        catalog=catalogDir+".csv"
        splitFile2RaMag(catalog, rootDir, datasetName)
        closeFiles(files)
        files={}
    ##
###########
testSplit()
###########
end_time = datetime.now()
print('{0} End'.format(end_time))
print('\nDone! Time elapsed {0}'.format(end_time-start_time))
